package edu.zju.gis.dbfg.model.index;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.XZQUtil;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;
import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_PYRAMID_CONFIG;

public class SHDIndex {

    private final static Logger logger = LoggerFactory.getLogger(SHDIndex.class);

    public static void main(String[] args) throws Exception{
        //Setup Environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setMaster("local[4]");
        conf.setAppName("SHDI");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup Model Arguments
        logger.info("Setup arguments");
        SHDIndexArgs mArgs = SHDIndexArgs.initArgs(args);
        if(mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(SHDIndex.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs()));
        CoordinateReferenceSystem targetCrs = DEFAULT_CRS();
        PyramidConfig targetPyramidConfig = DEFAULT_PYRAMID_CONFIG();

        MathTransform transform = null;
        if(!crs.equals(targetCrs)){
            transform = CRS.findMathTransform(crs, targetCrs);
        }

        //Read data from datasource and map it to geoItem
        String input = mArgs.getInput();
        String type = mArgs.getType();
        String dlbm = mArgs.getDlbmIndex();
        String area = mArgs.getAreaIndex();
        String zldwdm = mArgs.getZldwdmIndex();
        String id = mArgs.getIdIndex();
        int level = XZQUtil.getXZQLevel(mArgs.getLevel());
        Map<String, Object> attrs = new HashMap<>();
        attrs.put("dlbm",dlbm);
        attrs.put("area", area);
        attrs.put("zldwdm", zldwdm);
        JavaRDD<GeoItem> gi = null;
        if(type.equals("ES")){
            gi = DataReaderRDD.readToRDD2(type, jsc, input, attrs, null, null, transform);
        }
        else{
            gi = DataReaderRDD.readToRDD(type,jsc, input, id, attrs, transform);
        }

        Broadcast<Integer> levelbroad = jsc.broadcast(level);

        //获取面积值，并按照 <行政区-一级地类，面积> 的形式组织
        JavaPairRDD<String, Double> zldlRDD = gi.mapToPair(new PairFunction<GeoItem, String, Double>() {
            @Override
            public Tuple2<String, Double> call(GeoItem input) throws Exception {
                String sdlbm = String.valueOf(input.getProp().get("dlbm")).substring(0,2);
                String szldwdm = String.valueOf(input.getProp().get("zldwdm")).substring(0,levelbroad.getValue());
                Double darea = Double.valueOf(String.valueOf(input.getProp().get("area")));
                String key = String.format("%s-%s", szldwdm, sdlbm);
                return new Tuple2<String, Double>(key, darea);
            }
        }).reduceByKey(new Function2<Double, Double, Double>() {
            @Override
            public Double call(Double input1, Double input2) throws Exception {
                return input1 + input2;
            }
        });
        zldlRDD.cache();
        List<Tuple2<String, Double>> temp = zldlRDD.collect();

        //统计行政区划总面积
        JavaPairRDD<String, Double> zlRDD = zldlRDD.mapToPair(new PairFunction<Tuple2<String, Double>, String, Double>() {
            @Override
            public Tuple2<String, Double> call(Tuple2<String, Double> input) throws Exception {
                String key = input._1.split("-")[0];
                return new Tuple2<>(key, input._2);
            }
        }).reduceByKey((x1,x2)->x1+x2);

        zlRDD.cache();
        List<Tuple2<String, Double>> temp2 = zlRDD.collect();

        //以行政区划为key组织各类面积数据
        JavaPairRDD<String, Tuple2<String, Double>> zdRDD = zldlRDD.mapToPair(new PairFunction<Tuple2<String, Double>, String, Tuple2<String, Double>>() {
            @Override
            public Tuple2<String, Tuple2<String, Double>> call(Tuple2<String, Double> input) throws Exception {
                String key = input._1.split("-")[0];
                Tuple2<String, Double> value = new Tuple2<>(input._1.split("-")[1], input._2);
                return new Tuple2<>(key, value);
            }
        });

        zdRDD.cache();
        List<Tuple2<String, Tuple2<String, Double>>> temp3 = zdRDD.collect();

        //join以行政区划为key的 各类面积和总面积，计算比例
        JavaPairRDD<String, Double> resultRDD = zlRDD.join(zdRDD).mapToPair(new PairFunction<Tuple2<String, Tuple2<Double, Tuple2<String, Double>>>, String, Double>() {
            @Override
            public Tuple2<String, Double> call(Tuple2<String, Tuple2<Double, Tuple2<String, Double>>> input) throws Exception {
                Double darea = input._2._2._2();
                Double sumarea = input._2._1();
                Double ratio = darea/sumarea;
                Double lnratio = Math.log(ratio);
                Double value = ratio * lnratio;
                String key = input._1;
                return new Tuple2<>(key, value);
            }
        }).reduceByKey(new Function2<Double, Double, Double>() {  //按照SHDI公式计算
            @Override
            public Double call(Double input1, Double input2) throws Exception {
                return input1+input2;
            }
        });

        List<Tuple2<String, Double>> result = resultRDD.collect();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        writer.write(String.format("%s,%s\n", zldwdm, "SHDI"));
        for (Tuple2<String, Double> item : result) {
            String s = String.format("%s,%.9f", item._1, item._2);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();
    }
}
